package com.jsean.demo;

import org.rocksdb.*;

import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Properties;

public class RocksDBDemo {

    static {
        // 加载 RocksDB 库
        RocksDB.loadLibrary();
    }

    private static String getRocksDBDataDir() {
        Properties props = new Properties();
        try (InputStream is = RocksDBDemo.class.getClassLoader().getResourceAsStream("application.properties")) {
            props.load(is);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        return props.getProperty("app.data.dir", "rocksdb-data");
    }

    public static void main(String[] args) {
        // RocksDB 数据库路径
        String dbPath = getRocksDBDataDir();

        try (final Options options = new Options()
                .setCreateIfMissing(true)
                .setMaxOpenFiles(5000)
                .setWriteBufferSize(64 * 1024 * 1024)  // 64MB
                .setMaxWriteBufferNumber(3)
                .setMaxBackgroundCompactions(10)
                .setCompressionType(CompressionType.SNAPPY_COMPRESSION);) {
            // 打开数据库
            try (final RocksDB db = RocksDB.open(options, dbPath)) {
                System.out.println("RocksDB opened successfully");

                // 1. 写入数据
                byte[] key1 = "key1".getBytes();
                byte[] value1 = "value1".getBytes();
                db.put(key1, value1);

                byte[] key2 = "key2".getBytes();
                byte[] value2 = "value2".getBytes();
                db.put(key2, value2);

                System.out.println("写入数据完成");

                // 2. 读取数据
                byte[] getValue1 = db.get(key1);
                System.out.println("key1: " + new String(getValue1));

                byte[] getValue2 = db.get(key2);
                System.out.println("key2: " + new String(getValue2));

                // 3. 更新数据
                byte[] newValue1 = "new_value1".getBytes();
                db.put(key1, newValue1);
                System.out.println("更新后的key1: " + new String(db.get(key1)));

                // 4. 删除数据
                db.delete(key2);
                System.out.println("删除key2后是否存在: " + (db.get(key2) == null));

                // 5. 批量操作
                try (final org.rocksdb.WriteBatch batch = new org.rocksdb.WriteBatch()) {
                    batch.put("batch_key1".getBytes(), "batch_value1".getBytes());
                    batch.put("batch_key2".getBytes(), "batch_value2".getBytes());
                    batch.delete(key1);
                    db.write(new org.rocksdb.WriteOptions(), batch);
                }

                // 6. 遍历所有键值对
                System.out.println("\n当前数据库内容:");
                try (final RocksIterator iterator = db.newIterator()) {
                    for (iterator.seekToFirst(); iterator.isValid(); iterator.next()) {
                        System.out.println(
                                "Key: " + new String(iterator.key()) +
                                        ", Value: " + new String(iterator.value()));
                    }
                }

                // 7. 获取所有键
                System.out.println("\n所有键:");
                List<String> keys = new ArrayList<>();
                try (final RocksIterator iterator = db.newIterator()) {
                    for (iterator.seekToFirst(); iterator.isValid(); iterator.next()) {
                        keys.add(new String(iterator.key()));
                    }
                }
                keys.forEach(System.out::println);

            } catch (RocksDBException e) {
                System.err.println("操作RocksDB出错: " + e.getMessage());
            }
        } catch (Exception e) {
            System.err.println("初始化RocksDB选项出错: " + e.getMessage());
        }
    }
}